from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn import metrics
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import KFold
from assess_model import show_metrics
import warnings
import time
warnings.filterwarnings("ignore")

LABEL_NAMES = ["video", "music", "picture", "file", "video and voice call over ip", "Others"]

# 创建误差矩阵
def create_confusion_matrix(y_pred, y_test):
    # calculate the confusion matrix
    confmat = metrics.confusion_matrix(y_true=y_test, y_pred=y_pred)

    fig, ax = plt.subplots(figsize=(7, 7))
    ax.imshow(confmat, cmap=plt.cm.Blues, alpha=0.5)

    n_labels = len(LABEL_NAMES)
    ax.set_xticks(np.arange(n_labels))
    ax.set_yticks(np.arange(n_labels))
    ax.set_xticklabels(LABEL_NAMES)
    ax.set_yticklabels(LABEL_NAMES)

    # rotate the tick labels and set their alignment.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
    # loop over data dimensions and create text annotations.
    for i in range(confmat.shape[0]):
        for j in range(confmat.shape[1]):
            ax.text(x=i, y=j, s=confmat[i, j], va='center', ha='center')

    # avoid that the first and last row cut in half
    bottom, top = ax.get_ylim()
    ax.set_ylim(bottom + 0.5, top - 0.5)

    ax.set_title("Confusion Matrix")
    ax.set_xlabel('True Label')
    ax.set_ylabel('Predicted Label')

    plt.tight_layout()
    plt.show()


data = pd.read_csv("data_model.csv")
fs_name = ['Fwd Packet Length Mean','Fwd Packet Length Max','Packet Length Mean','Fwd Packet Length Std','Bwd Iat Min','Packet Length Std','Bwd Packet Length Std','Packet Length Max','Packet Iat Min','Fwd Flow Byte/s']
X = data[fs_name]
y = pd.read_csv("new_label.csv").values.ravel() 
scaler = StandardScaler()
X = scaler.fit_transform(X)

kf = KFold(n_splits = 10, shuffle=True, random_state=5)    # 10折

predicted_y = []
expected_y = []

start_time = time.time()
for train_index, test_index in kf.split(X,y):     # 将数据划分为k折
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

    # KNN分类器
    clf = KNeighborsClassifier(n_neighbors=6)
    # 拟合数据集
    clf = clf.fit(X_train, y_train)

    y_pred = clf.predict(X_test)
    
    # store result from classification
    predicted_y.extend(y_pred)
    # store expected result for this specific fold
    expected_y.extend(y_test)
end_time = time.time()
print("using {} s".format(end_time-start_time))
show_metrics(expected_y,predicted_y)

create_confusion_matrix(expected_y,predicted_y)

